import argparse
from PIL import Image
from torch.utils.data.dataset import ConcatDataset, Dataset
from utils.consts import DATA_PATH, DATA_CONF, ORIGIN_DATA_PATH
from utils.tools import check_path, write_json
from dataproc.optset import get_datasets


r'''
img, label = dataset[0]
assert img isinstance of PIL.Image.Image
assert label is int and in range(0, num_cls)
label_dict[i] = "cls_name" means label i name is "cls_name"
example:
cifar_10's label is in [0,1,2,3,4,5,6,7,8,9]
label_dict = ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"] 
'''


def get_distr(dataset: Dataset, num_cls: int):
    bucket = [0] * num_cls
    for _, label in dataset:
        bucket[label] += 1
    return bucket


def trans_data(name: str, dataset: Dataset, label_dict: list[str], vl_rate: float):
    # mkdir
    save_path = f"{DATA_PATH}/{name}"
    check_path(save_path, remove_old=True)
    for dir_name in ["train", "val"]:
        check_path(f"{save_path}/{dir_name}")
    for cls_name in label_dict:
        check_path(f"{save_path}/train/{cls_name}")
        check_path(f"{save_path}/val/{cls_name}")

    # split into tr and vl set
    num_cls = len(label_dict)
    distr = get_distr(dataset, num_cls)
    thresh = [round(x*vl_rate) for x in distr]
    bucket = [0] * num_cls
    for img, label in dataset:
        id, dir = None, None
        assert (isinstance(img, Image.Image))
        if bucket[label] >= thresh[label]:
            id, dir = f"{bucket[label]-thresh[label]}", 'train'
        else:
            id, dir = f"{bucket[label]}", 'val'
        img_path = f"{save_path}/{dir}/{label_dict[label]}/{id}.png"
        img.save(img_path)
        print(img_path)
        bucket[label] += 1
    record = {
        "label": label_dict,
        "val": thresh,
        "train": [distr[i]-thresh[i] for i in range(num_cls)]
    }
    write_json(f'{save_path}/{DATA_CONF}', record, 'a')


def trans_datas(name: str, tr_set: Dataset, vl_set: Dataset, label_dict: list[str], vl_rate: float = 0.16):
    all_set = ConcatDataset([tr_set, vl_set])
    trans_data(name, all_set, label_dict, vl_rate)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='origin data opt')
    parser.add_argument('--dataset', '-d', type=str, required=True)
    args = parser.parse_args()

    tr_set, vl_set, lb_dict = None, None, None
    name = args.dataset

    tr_set, vl_set, lb_dict = get_datasets(args.dataset, ORIGIN_DATA_PATH)

    trans_datas(name, tr_set, vl_set, lb_dict)
    print("Trans Success!\n")
